# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openFuyao is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
This module defines NetworkCollector class.
"""

import logging
import os
import re
from typing import Any

from hardware_monitor.base.base_metrics_collector import BaseCollector
from hardware_monitor.base.registry import registry
from hardware_monitor.base.utils import add_error
from hardware_monitor.collector.metrics.npu.npu_utils import init_device_info, get_phy_id, get_info_from_hccn_tool
from hardware_monitor.collector.metrics.npu.constants import RET_INVALID_VALUE

module_name = os.getenv("MODULE_NAME", "default")

logger = logging.getLogger(module_name)


class NetworkCollector(BaseCollector):
    """
    This collector collects information about the network.
    """
    @property
    def name(self):
        return "network"

    def _collect_impl(self):
        results = {}
        node_name = os.getenv("NODE_NAME", "")
        logic_ids = self.cache.get("logic_id") or []
        card_ids = self.cache.get("card_id") or []
        device_ids = self.cache.get("device_id") or []
        for i, logic_id in enumerate(logic_ids):
            card_id = card_ids[i]
            device_id = device_ids[i]
            info = self._get_device_all_network_info(node_name, logic_id, card_id, device_id)
            if info:
                results[str(logic_id)] = info
        return results

    @staticmethod
    def _get_device_all_network_info(node_name: str, logic_id: int, card_id: int, device_id: int) -> dict[str, Any]:
        fields = [
            "node_name", "card_id", "device_id", "phy_id",
            "npu_chip_info_bandwidth_tx", "npu_chip_info_bandwidth_rx", "npu_chip_info_link_status"
        ]

        result = init_device_info(node_name, card_id, device_id, fields)

        phy_id, err = get_phy_id(logic_id)
        if err is not None:
            msg = f"Failed to get physical ID from logicID({logic_id}): {err}"
            logger.error("[NetworkCollector] %s",msg)
            return {"error_code": phy_id, "error_message": msg}

        result.update({"phy_id": phy_id})

        npu_traffic = NetworkCollector.get_npu_interface_traffic(logic_id, phy_id)
        if "error_code" not in npu_traffic:
            result.update({
                "npu_chip_info_bandwidth_tx": npu_traffic["tx"],
                "npu_chip_info_bandwidth_rx": npu_traffic["rx"]})
        else:
            add_error(result, "network_info", npu_traffic["error_code"], npu_traffic["error_message"])

        npu_link_status = NetworkCollector.get_npu_link_status(logic_id, phy_id)
        if "error_code" not in npu_link_status:
            result.update({
                "npu_chip_info_bandwidth_tx": npu_link_status["status"]})
        else:
            add_error(result, "network_info", npu_link_status["error_code"], npu_link_status["error_message"])

        return result

    @staticmethod
    def get_npu_interface_traffic(logic_id: int, phy_id: int) -> dict[str, Any]:
        """
        Retrieves the network interface traffic (TX/RX bandwidth)
        for a given NPU logic ID and physical ID.
        """
        args = ["-i", str(int(phy_id)), "-bandwidth", "-g"]
        out, ret_code, err = get_info_from_hccn_tool(*args)
        if err is not None:
            msg = f"Failed to get hccn_tool command exec result from logicID({logic_id}): {err}"
            logger.error("[NetworkCollector] %s",msg)
            return {"error_code": ret_code, "error_message": msg}

        bandwidth_pat = re.compile(
            r"Bandwidth\s+(TX|RX):\s*([0-9]*\.?[0-9]+)\s*MB/sec", re.IGNORECASE
        )

        tx = float(RET_INVALID_VALUE)
        rx = float(RET_INVALID_VALUE)

        for line in out.splitlines():
            m = bandwidth_pat.search(line)
            if not m:
                continue
            kind = m.group(1).upper()
            try:
                val = float(m.group(2))
            except ValueError:
                continue
            if kind == "TX":
                tx = val
            elif kind == "RX":
                rx = val

        if tx == float(RET_INVALID_VALUE) or rx == float(RET_INVALID_VALUE):
            msg = f"TX/RX traffic not fully available for logicID({logic_id}), tx={tx}, rx={rx}"
            logger.warning("[NetworkCollector] %s",msg)
            return {"error_code": RET_INVALID_VALUE, "error_message": msg}

        return {
            "tx": tx,
            "rx": rx
        }

    @staticmethod
    def get_npu_link_status(logic_id: int, phy_id: int) -> dict[str, Any]:
        """
        Retrieves the link status for a given NPU logic ID and physical ID.
        """
        args = ["-i", str(int(phy_id)), "-link", "-g"]
        out, ret_code, err = get_info_from_hccn_tool(*args)
        if err is not None:
            msg = f"Failed to get hccn_tool command exec result from logicID({logic_id}): {err}"
            logger.error("[NetworkCollector] %s",msg)
            return {"error_code": ret_code, "error_message": msg}

        replaced = out.replace("\r", "").replace("\n", "")
        parts = replaced.split()
        if len(parts) == 3 and parts[0].lower() == "link" and parts[1].lower() == "status:":
            status = parts[2].upper()
            return {"status": status}
        msg = f"Failed to get npu link status from logicID({logic_id})"
        logger.error("[NetworkCollector] %s",msg)
        return {"error_code": RET_INVALID_VALUE, "error_message": msg}


collector = NetworkCollector()
registry.register(collector.name, collector, group="collector.npu")
